import torch
from torchvision import datasets, transforms
import matplotlib.pyplot as plt
import torchvision
import os
import sys
import time
import glob
import numpy as np
import torch
import utils
import logging
import argparse
import torch.nn as nn
import torch.utils
import torch.nn.functional as F
import torchvision.datasets as dset
import torch.backends.cudnn as cudnn
# 设置数据转换
transform = transforms.Compose([
    transforms.Resize((256, 256)),  # 根据需要调整图像尺寸
    transforms.ToTensor()
])

# 加载数据集
dataset = datasets.ImageFolder(root='D:\\ai\\code\\data\\archive\\train_valid_test\\train', transform=transform)

# 创建数据加载器
dataloader = torch.utils.data.DataLoader(dataset, batch_size=4, shuffle=True)

# 获取一批图像
images, labels = next(iter(dataloader))

# 反转换 tensor，以便可视化
def imshow(img):
    img = img / 2 + 0.5  # 反归一化
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()

# 显示图像和标签
print('Labels:', [dataset.classes[label] for label in labels])
imshow(torchvision.utils.make_grid(images))
